package zk.learn.graphx

import org.apache.spark.graphx.{Edge, Graph, VertexId}
import org.apache.spark.rdd.RDD
import org.apache.spark.{SparkConf, SparkContext}

//noinspection DuplicatedCode
object Airport {
  def main(args: Array[String]): Unit = {
    val conf: SparkConf = new SparkConf().setAppName(this.getClass.getCanonicalName)
      .setMaster("local[*]")
    val sc = new SparkContext(conf)
    sc.setLogLevel("warn")

    // 定义顶点
    val vertexArray: Array[(VertexId, String)] = Array(
      (1L, "SFO"),
      (2L, "ORD"),
      (3L, "DFW"))
    val vertexRDD: RDD[(VertexId, String)] = sc.makeRDD(vertexArray)

    // 定义边
    val edgeArray: Array[Edge[Int]] = Array(
      Edge(1L, 2L, 1800),
      Edge(2L, 3L, 800),
      Edge(3L, 1L, 1400),
    )
    val edgeRDD: RDD[Edge[Int]] = sc.makeRDD(edgeArray)

    // 图的定义
    val graph: Graph[String, Int] = Graph(vertexRDD, edgeRDD)


    println("\n\n\n求所有的顶点:")
    graph.vertices
        .foreach(println)

    println("\n\n\n求所有的边:")
    graph.edges
      .foreach(println)

    println("\n\n\n求所有的triplets:")
    graph.triplets
      .foreach(println)

    println(s"\n\n\n求顶点数:${graph.numVertices}")

    println(s"\n\n\n求边数:${graph.numEdges}")

    println("\n\n\n求机场距离大于1000的有几个，有哪些按所有机场之间的距离排序（降序）")

    // 求指定顶点到图中其他顶点到最短距离, 返回结果类型是RDD[(from:VertexId, to:VertexId, dist:Double)
    val minDists1 = minDists(graph, 1L)
    val minDists2 = minDists(graph, 2L)
    val minDists3 = minDists(graph, 3L)

    minDists1
      .union(minDists2).union(minDists3) // 合并三个最短距离的集合
        .filter{
          case (_, _, dist) => dist > 1000
        }// dist 需要大于1000
      .sortBy(_._3, ascending = false)// 按照dist排序
      .collect()
      .foreach(info =>
        {
          println(s"顶点${info._1}到顶点${info._2}到距离为${info._3}")
        }
      )// 输出
    //注： 需要先collect， 如果直接foreach, 将会是乱序的


    sc.stop()
  }



  def minDists(graph: Graph[String, Int], sourceId:VertexId):RDD[(VertexId, VertexId, Double)] = {
    val initailGraph: Graph[Double, Int] = graph.mapVertices((id, _) => if (id == sourceId) 0.0 else Double.PositiveInfinity)

    val disGraph: Graph[Double, Int] = initailGraph.pregel(Double.PositiveInfinity)(
      // 两个消息来的时候，取其中的最小路径
      (_, dist, newDist) => math.min(dist, newDist),

      // Send Message 函数
      triplet => {
        if (triplet.srcAttr + triplet.attr < triplet.dstAttr) {
          Iterator((triplet.dstId, triplet.srcAttr + triplet.attr))
        } else
          Iterator.empty
      },

      // mergeMsg
      (dista, distb) => math.min(dista, distb)
    )

    // 转换一下数据类型
    disGraph.vertices.map{
      case (targetId, dist) =>
        (sourceId, targetId, dist)
    }
  }
}

